package com.jvyou.app.ds.client;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jvyou.app.ds.domain.dto.SparkRequestDto;
import com.jvyou.app.ds.domain.dto.SparkResponseDto;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.List;

/**
 * 讯飞星火认知大模型客户端
 *
 * @author 橘柚
 * @since 2024/12/20 12:13
 */

@Slf4j
@Service
public class SparkServiceClient {

    private final ObjectMapper objectMapper = new ObjectMapper();

    public SparkResponseDto chat(SparkRequestDto sparkRequestDto) {
        log.info("正在连接generalv3.5大模型......");
        RestTemplate restTemplate = new RestTemplate();

        String url = "https://spark-api-open.xf-yun.com/v1/chat/completions";
        log.info("generalv3.5请求地址：{}", url);
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer fTMZcgpapAbZetwizBSn:UZedeOMeLZrcxKXtUqCa");
        headers.setContentType(MediaType.APPLICATION_JSON);

        String requestBody;
        try {
            log.info("开始构建请求体.....");
            requestBody = objectMapper.writeValueAsString(sparkRequestDto);
            log.info("构建请求体成功:{}", requestBody);
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        }
        HttpEntity<String> request = new HttpEntity<>(requestBody, headers);
        ResponseEntity<String> response = restTemplate.exchange(url, HttpMethod.POST, request, String.class);
        log.info("星火大模型响应结果:{}", response.getBody());
        return parseSparkResponse(response.getBody());
    }


    public static SparkResponseDto parseSparkResponse(String jsonStr) {
        // 使用Fastjson解析JSON字符串
        JSONObject jsonObject = JSON.parseObject(jsonStr);

        // 创建响应对象
        SparkResponseDto response = new SparkResponseDto();
        response.setCode(jsonObject.getInteger("code"));
        response.setMessage(jsonObject.getString("message"));
        response.setSid(jsonObject.getString("sid"));

        // 解析choices数组
        List<SparkResponseDto.Choice> choices = JSON.parseArray(
                jsonObject.getJSONArray("choices").toJSONString(),
                SparkResponseDto.Choice.class
        );
        response.setChoices(choices);

        // 解析usage对象
        JSONObject usageJson = jsonObject.getJSONObject("usage");
        SparkResponseDto.Usage usage = new SparkResponseDto.Usage();
        usage.setPromptTokens(usageJson.getInteger("prompt_tokens"));
        usage.setCompletionTokens(usageJson.getInteger("completion_tokens"));
        usage.setTotalTokens(usageJson.getInteger("total_tokens"));
        response.setUsage(usage);
        return response;
    }

}
